/*
 * Copyright 2012 David Crosson
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package fr.janalyse.jmx

import scala.util.{Properties=>SP}
import java.io.File.{separator=>FS, pathSeparator=>PS}
import scala.reflect.BeanProperty
import java.lang.management.ManagementFactory
import javax.management.ObjectName
import javax.management.ObjectInstance
import javax.management.ObjectName
import javax.management.remote.rmi.RMIConnection
import javax.management.remote.rmi.RMIServer
import javax.management.remote.rmi.RMIServerImpl_Stub
import javax.management.openmbean.{CompositeData, CompositeDataSupport}
import javax.management.openmbean.{TabularData,TabularDataSupport}
import scala.collection.JavaConversions._
import java.rmi.MarshalledObject
import java.rmi.UnmarshalException
import javax.management.openmbean.CompositeData
import scala.collection.JavaConversions._
import javax.management.MBeanAttributeInfo
import java.rmi.registry.LocateRegistry
import javax.management.remote.JMXServiceURL
import javax.management.remote.JMXConnectorServerFactory
import java.rmi.server.RMIClientSocketFactory
import javax.management.remote.rmi.RMIJRMPServerImpl
import java.rmi.server.RMIServerSocketFactory
import javax.net.ServerSocketFactory
import java.net.ServerSocket
import java.rmi.server.RMISocketFactory
import javax.management.RuntimeMBeanException


trait JMXAutoClose {
  def usingJMX[T <: { def close() }, R](resource: T)(block: T => R) = {
    try block(resource)
    finally resource.close
  }
}


case class CompositeDataWrapper(cd:CompositeData) {
  // TODO : BAD because it assumes that all values are of the same type...
  def toMap[RT]() = {
    val keys = cd.getCompositeType().keySet()
    val tuples = keys.toList map {k => (k -> cd.get(k).asInstanceOf[RT]) }
    tuples.toMap
  }
}

case class CompositeDataSupportWrapper(cd:CompositeDataSupport) {
  // TODO : BAD because it assumes that all values are of the same type...
  def toMap[RT]() = {
    val keys = cd.getCompositeType().keySet()
    val tuples = keys.toList map {k => (k -> cd.get(k).asInstanceOf[RT]) }
    tuples.toMap
  }
}

case class TabularDataWrapper(tabularData:TabularData) {
  def toMap[X]() = {
      val indexNamesInCell = tabularData.getTabularType().getIndexNames()
	  val tuples = tabularData.values() collect {
		case (k:java.util.List[String], v:CompositeData) => {
			val cellContent = CompositeDataWrapper(v).toMap[X]
			val key    = indexNamesInCell map {cellContent.get(_).get} mkString("-")
			val submap = cellContent filterNot {case (k,_) => indexNamesInCell contains k}
			
			key -> submap
		}
	  }
      tuples.toMap
  }
}


case class TabularDataSupportWrapper(tabularData:TabularDataSupport) {
  def toMap[X]() = {
      val indexNamesInCell = tabularData.getTabularType().getIndexNames()
	  val tuples = tabularData collect {
		case (k:java.util.List[String], v:CompositeDataSupport) => {
			val cellContent = CompositeDataSupportWrapper(v).toMap[X]
			val key    = indexNamesInCell map {cellContent.get(_).get} mkString("-")
			val submap = cellContent filterNot {case (k,_) => indexNamesInCell contains k}
			
			key -> submap
		}
	  }
      tuples.toMap
  }
}





object RichMBean  {
  def apply(name:String)(implicit jmx:JMX) = new RichMBean(name)
}


 
case class RichMBean(objectName:ObjectName)(implicit jmx:JMX) {
  
  implicit private lazy val objectInstance = jmx.rmiConnection.getObjectInstance(objectName, null) match {
    case null  => throw new RuntimeException("Couldn't find MBean")
    case mbean => mbean
  }
  val name = objectName.toString()
  //val keys = mapAsScalaMap(objectName.getKeyPropertyList())  // Not available in scala 2.8.x
  val keys = (objectName.getKeyPropertyList().entrySet().toList map {case entry => entry.getKey() -> entry.getValue()}).toMap
  
  def mbean = this
  def getOption[A](name:String):Option[A] = {
    val mbeanInfo = jmx.rmiConnection.getMBeanInfo(objectName, null)
    val attrsInfo:List[javax.management.MBeanAttributeInfo] = mbeanInfo.getAttributes().toList
    (attrsInfo find {attrInfo => attrInfo.getName == name}) match {
      case None => None
      case Some(attrInfo) => get[A](name) match {
          case x if (x==null) => None
          case x => Some(x)
      }
    }
  }
  def get[A](attrname:String) = jmx.rmiConnection.getAttribute(objectName,attrname,null).asInstanceOf[A]
  def set[A](attrname:String, value:A) {
    val attribute = new javax.management.Attribute(attrname, value)
    jmx.rmiConnection.setAttribute(objectName, attribute, null)
  }
  def get[A](attrnames:String*):List[A] = attrnames.toList map { get[A](_)}
  def call[A](operation:String, args:Any*) : Option[A] = {
    val signature=args map { // http://home.pacifier.com/~mmead/jni/cs510ajp/index.html  /   http://en.wikipedia.org/wiki/Java_Native_Interface#Mapping_types
      case _:Int     => "int"
      case _:Short   => "short"
      case _:Long    => "long"
      case _:Float   => "float"
      case _:Double  => "double"
      case _:Boolean => "boolean"
      case _:Byte    => "byte"
      case _:Char    => "char"
      case _:String  => "String"
      case _:Array[Int]     => "[I"
      case _:Array[Short]   => "[S"
      case _:Array[Long]    => "[J"
      case _:Array[Float]   => "[F"
      case _:Array[Double]  => "[D"
      case _:Array[Boolean] => "[Z"
      case _:Array[Byte]    => "[B"
      case _:Array[Char]    => "[C"
      case _:Array[String]  => "[Ljava/lang/String"
    }
    val args2give=args.toArray
    jmx.rmiConnection.invoke(objectName, operation, args2give, signature.toArray, null) match {
      case null => None
      case result => Some(result.asInstanceOf[A])
    }
  }

  
  private def genericGetter[T, R<:RichAttribute](attr:R, getter:(Object)=>T):Option[T] = {
    try {
	    jmx.rmiConnection.getAttribute(objectName,attr.name,null) match {
	      case null => None
	      case x => Some(getter(x))
	    }
    } catch {
      case e:RuntimeMBeanException if e.getCause().isInstanceOf[UnsupportedOperationException] => None
      case e:UnmarshalException => None
      case x =>
          //println("Warning: Error while getting value for attribute %s mbean %s (%s)".format(attr.name, name, x))
        None
    }    
  }
  
  def getString(attr:RichAttribute):Option[String] = genericGetter(attr, attr.asString)
  def getDouble(attr:RichNumberAttribute):Option[Double] = genericGetter(attr, attr.asDouble)
  def getLong(attr:RichNumberAttribute):Option[Long] = genericGetter(attr, attr.asLong)
  
  def attributes():List[RichAttribute] = {
    val mbeanInfo = jmx.rmiConnection.getMBeanInfo(objectName, null)
    val attrsInfo:List[MBeanAttributeInfo] = mbeanInfo.getAttributes().toList
    attrsInfo map {ai =>
      val n=ai.getName()
      val d= if (ai.getDescription()==ai.getName) None else Some(ai.getDescription())
      ai.getType() match {
        case "java.lang.Boolean"|"boolean"|"java.lang.boolean" => RichBooleanAttribute(n,d)
        case "java.lang.Byte"|"byte"       => RichByteAttribute(n,d)
        case "java.lang.Short"|"short"     => RichShortAttribute(n,d)
        case "java.lang.Integer"|"int"     => RichIntAttribute(n,d)
        case "java.lang.Long"|"long"       => RichLongAttribute(n,d)
        case "java.lang.Float"|"float"     => RichFloatAttribute(n,d)
        case "java.lang.Double"|"double"   => RichDoubleAttribute(n,d)
        case "java.lang.String"|"String"   => RichStringAttribute(n,d)
        //case "java.util.List" =>
        //case "java.util.Properties" =>
        case "[Ljava.lang.String;"  => RichStringArrayAttribute(n,d)
        //case "javax.management.openmbean.TabularData"      =>
        case "javax.management.openmbean.CompositeData"    => RichCompositeDataAttribute(n,d)
        //case "[Ljavax.management.openmbean.CompositeData;" =>
        //case "[Ljavax.management.ObjectName;" => 
        case x =>
          //println("Warning: Not supported jmx attribute value type %s for %s mbean %s".format(x, n, name))
          //println(x)
          RichGenericAttribute(ai.getName(),d)
      }
    }
  //def apply[A](attrname:String) = jmx.rmiConnection.getAttribute(objectName,attrname,null).asInstanceOf[A]
  }
}


/*
import fr.janalyse.jmx._
val jmx=JMX()
val os=jmx.os.get
val mem=jmx.memory.get
val th=jmx.threading.get
 */

trait RichAttribute {
  val name:String
  val desc:Option[String]
  def asString(ob:Object):String = ob.toString()
}

trait RichArrayAttribute extends RichAttribute {
}

trait RichNumberAttribute extends RichAttribute {
  def asDouble(ob:Object):Double = ob match {
    case e:java.lang.Byte => e.toDouble
    case e:java.lang.Short => e.toDouble
    case e:java.lang.Integer => e.toDouble
    case e:java.lang.Long => e.toDouble
    case e:java.lang.Float => e.toDouble
    case e:java.lang.Double => e
    case e:java.lang.String => e.toDouble
  }
  def asLong(ob:Object):Long = ob match {
    case e:java.lang.Byte => e.toLong
    case e:java.lang.Short => e.toLong
    case e:java.lang.Integer => e.toLong
    case e:java.lang.Long => e
    case e:java.lang.Float => e.toLong
    case e:java.lang.Double => e.toLong
    case e:java.lang.String => e.toLong
  }

}

case class RichCompositeDataAttribute(name:String, desc:Option[String]=None) extends RichAttribute {
  override def asString(ob:Object):String = {
    val cd = ob.asInstanceOf[CompositeData]
    cd.toString()
  }
}

case class RichStringArrayAttribute(name:String, desc:Option[String]=None) extends RichArrayAttribute {
  override def asString(ob:Object):String = {
    ob.asInstanceOf[Array[String]].toList.mkString(", ")
  }
}

case class RichBooleanAttribute(name:String, desc:Option[String]=None) extends RichAttribute

case class RichStringAttribute(name:String, desc:Option[String]=None) extends RichAttribute

case class RichGenericAttribute(name:String, desc:Option[String]=None) extends RichAttribute


case class RichByteAttribute(name:String, desc:Option[String]=None) extends RichNumberAttribute
case class RichShortAttribute(name:String, desc:Option[String]=None) extends RichNumberAttribute
case class RichIntAttribute(name:String, desc:Option[String]=None) extends RichNumberAttribute
case class RichLongAttribute(name:String, desc:Option[String]=None) extends RichNumberAttribute
case class RichDoubleAttribute(name:String, desc:Option[String]=None) extends RichNumberAttribute
case class RichFloatAttribute(name:String, desc:Option[String]=None) extends RichNumberAttribute



case class Credentials(username:String="", password:String="")

case class JMXOptions(
  host: String = "localhost",
  port: Int = 1099,
  name: Option[String] = None,
  username: Option[String] = None,
  password: Option[String] = None,
  connectTimeout: Long = 30000,
  retryCount: Int = 5,
  retryDelay: Long = 2000
  ) {
  val credentials:Option[Credentials] = (username,password) match {
    case (None,None) => None
    case (None,_) => None
    case (Some(u), Some(p)) => Some(Credentials(u,p))
  }
}



object JMX extends JMXAutoClose {
  def register(ob:Object, obname:ObjectName) = ManagementFactory.getPlatformMBeanServer.registerMBean(ob, obname)
  
  def once[R](
      host:String,
      port:Int,
      username:Option[String]=None,
      password:Option[String]=None) (block: JMX => R):R = {
    once[R](JMXOptions(host=host, port=port, username=username, password=password))(block)
  }
  
  def once[R]() (block: JMX => R) = once[R](None)(block)    
  def once[R](options:JMXOptions) (block: JMX => R) = once[R](Some(options))(block)    
  def once[R](someOptions:Option[JMXOptions]) (block: JMX => R) = usingJMX(new JMX(someOptions)) { block(_)} 

    
  def apply(options:JMXOptions) = new JMX(Some(options))
  def apply(options:Option[JMXOptions]) = new JMX(options)
  def apply(host:String, port:Int) = new JMX(Some(JMXOptions(host=host, port=port)))
  def apply() = new JMX()
}





class JMX(val options:Option[JMXOptions]=None) {
  //def jmx=this
  lazy val rmiConnection = options map {o =>
    findMBeanServers(o.host, o.port, o.credentials).head
  } getOrElse {
    // By default - self connect...
    // TODO : BAD BAD BAD - Find a better way...
    val csf=RMISocketFactory.getDefaultSocketFactory()
    var chosenPort:Option[Int]=None
    val ssf = new RMIServerSocketFactory {
      def createServerSocket(port:Int):ServerSocket = {
        val ssock=ServerSocketFactory.getDefault().createServerSocket(port)
        chosenPort=Some(ssock.getLocalPort())
        ssock
      }
    }
    LocateRegistry.createRegistry(0, csf, ssf)
    val url = new JMXServiceURL("service:jmx:rmi:///jndi/rmi://localhost:%d/jmxapitestrmi".format(chosenPort.get))
    val mbs = ManagementFactory.getPlatformMBeanServer()
    val cs = JMXConnectorServerFactory.newJMXConnectorServer(url, null, mbs)
    cs.start()
    Thread.sleep(500)
    findMBeanServers("localhost", chosenPort.get, None).head
  }
  
  private def findMBeanServers(host:String, port:Int, credentials:Option[Credentials]) :List[RMIConnection] = {
    try {
      val lr = java.rmi.registry.LocateRegistry.getRegistry(host, port)
       //filterNot => some jonas MBeans causes registry data corruption if lookuped then all others call are in exception
      val entries = lr.list.toList
            .filterNot( _ contains "TMFactory" )
            .filterNot( _ contains "RMI_SERVER_RPC")
      entries map { entry =>
        try {
          val credentialsArray:Array[String] = credentials match {
            case None => Array()
            case Some(Credentials(username, password)) => Array(username, password)
          }
          lr.lookup(entry) match {
            case stub:RMIServerImpl_Stub => Some(stub.newClient(credentialsArray))
            case rsi:RMIServer => Some(rsi.newClient(credentialsArray))
            case r@_ => println("Do not know how to manage %s class %s".format(entry, r.getClass)); None
          }
        } catch {
          case e:UnmarshalException => /*println("%s - Missing = %s".format(entry, searchMissingClass(e)));*/  None
        }
      } flatten
    }
  }
  def close() = {
    try {
      rmiConnection.close
    } catch {
      case e => println("Exception while closing rmi connection, let's ignore and continue...")
    }
  }
  def searchMissingClass(e:Throwable):String = {
    e match {
      case null => "?"
      case ex:ClassNotFoundException => ex.getMessage().trim().split("\\s+")(0)
      case _ => searchMissingClass(e.getCause)
    }
  }
  def apply(name:String) = get(name)
  def get(name:String) = RichMBean(name)(this)
  def domains = rmiConnection.getDomains(null).toList
  def getOption(name:String) = if (rmiConnection.queryNames(name, null, null).size>0) Some(get(name)) else None
  //
  def mbeans(query:String) = {
    rmiConnection.queryNames(null,query,null).toList map {RichMBean(_)(this)}
  }
  def mbeans() = {
    rmiConnection.queryNames(null,null,null).toList map {RichMBean(_)(this)}
  }
  def os()        = getOption("java.lang:type=OperatingSystem")
  def runtime()   = getOption("java.lang:type=Runtime")
  def threading() = getOption("java.lang:type=Threading")
  def memory()    = getOption("java.lang:type=Memory")
  /*
  def systemProperties():Map[String,String] = {
    runtime().map { r => r.get[TabularDataSupport]("SystemProperties") collect {
    	case (k:java.util.List[String],v:CompositeDataSupport) 
    		if (v.values.size==2) => k.head.toString -> v.values.tail.head.toString 
      }
    }  getOrElse Map.empty
  }
  */
  def systemProperties():Map[String,String] = {
    runtime() match {
      case Some(runtime) =>
        runtime.get[TabularDataSupport]("SystemProperties") flatMap {
          case (k:java.util.List[String],v:CompositeDataSupport) if (v.values.size==2) =>
            k.head.toString -> v.values.tail.head.toString :: Nil   // values first item is the system property name !
          case _ => Nil
        } toMap
      case None => Map.empty
    }
  }

}

